// Copyright (c) 2016 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/tools/quic/chlo_extractor.h"

#include "net/quic/quic_framer.h"
#include "net/quic/test_tools/crypto_test_utils.h"
#include "net/quic/test_tools/quic_test_utils.h"

using base::StringPiece;
using std::string;
using testing::_;
using testing::Return;

namespace net {
namespace test {
    namespace {

        class TestDelegate : public ChloExtractor::Delegate {
        public:
            TestDelegate() { }
            ~TestDelegate() override { }

            // ChloExtractor::Delegate implementation
            void OnChlo(QuicVersion version,
                QuicConnectionId connection_id,
                const CryptoHandshakeMessage& chlo) override
            {
                version_ = version;
                connection_id_ = connection_id;
                chlo_ = chlo.DebugString();
            }

            QuicConnectionId connection_id() const { return connection_id_; }
            QuicVersion version() const { return version_; }
            const string& chlo() const { return chlo_; }

        private:
            QuicConnectionId connection_id_;
            QuicVersion version_;
            string chlo_;
        };

        class ChloExtractorTest : public ::testing::Test {
        public:
            ChloExtractorTest()
            {
                header_.public_header.connection_id = 42;
                header_.public_header.connection_id_length = PACKET_8BYTE_CONNECTION_ID;
                header_.public_header.version_flag = true;
                header_.public_header.versions = SupportedVersions(QuicSupportedVersions().front());
                header_.public_header.reset_flag = false;
                header_.public_header.packet_number_length = PACKET_6BYTE_PACKET_NUMBER;
                header_.packet_number = 1;
                header_.entropy_flag = false;
                header_.entropy_hash = 0;
            }

            void MakePacket(QuicStreamFrame* stream_frame)
            {
                QuicFrame frame(stream_frame);
                QuicFrames frames;
                frames.push_back(frame);
                QuicFramer framer(SupportedVersions(header_.public_header.versions.front()),
                    QuicTime::Zero(), Perspective::IS_CLIENT);
                std::unique_ptr<QuicPacket> packet(
                    BuildUnsizedDataPacket(&framer, header_, frames));
                EXPECT_TRUE(packet != nullptr);
                size_t encrypted_length = framer.EncryptPayload(
                    ENCRYPTION_NONE, header_.path_id, header_.packet_number, *packet,
                    buffer_, arraysize(buffer_));
                ASSERT_NE(0u, encrypted_length);
                packet_.reset(new QuicEncryptedPacket(buffer_, encrypted_length));
                EXPECT_TRUE(packet_ != nullptr);
                delete stream_frame;
            }

        protected:
            TestDelegate delegate_;
            QuicPacketHeader header_;
            std::unique_ptr<QuicEncryptedPacket> packet_;
            char buffer_[kMaxPacketSize];
        };

        TEST_F(ChloExtractorTest, FindsValidChlo)
        {
            CryptoHandshakeMessage client_hello;
            client_hello.set_tag(kCHLO);

            string client_hello_str(
                client_hello.GetSerialized().AsStringPiece().as_string());
            // Construct a CHLO with each supported version
            for (QuicVersion version : QuicSupportedVersions()) {
                QuicVersionVector versions(SupportedVersions(version));
                header_.public_header.versions = versions;
                MakePacket(
                    new QuicStreamFrame(kCryptoStreamId, false, 0, client_hello_str));
                EXPECT_TRUE(ChloExtractor::Extract(*packet_, versions, &delegate_))
                    << QuicVersionToString(version);
                EXPECT_EQ(version, delegate_.version());
                EXPECT_EQ(header_.public_header.connection_id, delegate_.connection_id());
                EXPECT_EQ(client_hello.DebugString(), delegate_.chlo())
                    << QuicVersionToString(version);
            }
        }

        TEST_F(ChloExtractorTest, DoesNotFindValidChloOnWrongStream)
        {
            CryptoHandshakeMessage client_hello;
            client_hello.set_tag(kCHLO);

            string client_hello_str(
                client_hello.GetSerialized().AsStringPiece().as_string());
            MakePacket(
                new QuicStreamFrame(kCryptoStreamId + 1, false, 0, client_hello_str));
            EXPECT_FALSE(
                ChloExtractor::Extract(*packet_, QuicSupportedVersions(), &delegate_));
        }

        TEST_F(ChloExtractorTest, DoesNotFindValidChloOnWrongOffset)
        {
            CryptoHandshakeMessage client_hello;
            client_hello.set_tag(kCHLO);

            string client_hello_str(
                client_hello.GetSerialized().AsStringPiece().as_string());
            MakePacket(new QuicStreamFrame(kCryptoStreamId, false, 1, client_hello_str));
            EXPECT_FALSE(
                ChloExtractor::Extract(*packet_, QuicSupportedVersions(), &delegate_));
        }

        TEST_F(ChloExtractorTest, DoesNotFindInvalidChlo)
        {
            MakePacket(new QuicStreamFrame(kCryptoStreamId, false, 0, "foo"));
            EXPECT_FALSE(
                ChloExtractor::Extract(*packet_, QuicSupportedVersions(), &delegate_));
        }

    } // namespace
} // namespace test
} // namespace net
